package com.autoglm.android.core.vision

import android.content.Context
import android.graphics.Bitmap
import android.graphics.Rect
import android.util.Log
import com.autoglm.android.core.llm.LLMService
import com.autoglm.android.core.models.UIElement
import com.autoglm.android.core.models.VisualElement
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.objects.ObjectDetection
import com.google.mlkit.vision.objects.defaults.ObjectDetectorOptions
import com.google.mlkit.vision.text.TextRecognition
import com.google.mlkit.vision.text.chinese.ChineseTextRecognizerOptions
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.coroutines.withContext
import org.json.JSONArray
import org.json.JSONObject
import kotlin.coroutines.resume
import kotlin.coroutines.resumeWithException

/**
 * 图像分析器，通过计算机视觉技术识别界面元素
 */
class ImageAnalyzer(
    private val context: Context,
    private val llmService: LLMService
) {
    companion object {
        private const val TAG = "ImageAnalyzer"
    }
    
    // 初始化文字识别器
    private val textRecognizer = TextRecognition.getClient(ChineseTextRecognizerOptions.Builder().build())
    
    // 初始化对象检测器
    private val objectDetector = ObjectDetection.getClient(
        ObjectDetectorOptions.Builder()
            .setDetectorMode(ObjectDetectorOptions.SINGLE_IMAGE_MODE)
            .enableMultipleObjects()
            .enableClassification()
            .build()
    )
    
    /**
     * 分析屏幕截图，识别界面元素
     * 
     * @param screenshot 屏幕截图
     * @return 识别到的视觉元素列表
     */
    suspend fun analyzeScreenshot(screenshot: Bitmap): List<VisualElement> = withContext(Dispatchers.Default) {
        val elements = mutableListOf<VisualElement>()
        
        try {
            // 识别文字
            val textElements = detectText(screenshot)
            elements.addAll(textElements)
            
            // 识别界面元素 (如按钮、输入框等)
            val uiElements = detectUIElements(screenshot)
            elements.addAll(uiElements)
            
            // 使用大模型分析界面更复杂的组件
            val complexElements = analyzeComplexElements(screenshot, elements)
            elements.addAll(complexElements)
            
            return@withContext elements
        } catch (e: Exception) {
            Log.e(TAG, "Error analyzing screenshot", e)
            return@withContext emptyList()
        }
    }
    
    /**
     * 检测图像中的文字
     */
    private suspend fun detectText(bitmap: Bitmap): List<VisualElement> = withContext(Dispatchers.Default) {
        val elements = mutableListOf<VisualElement>()
        val image = InputImage.fromBitmap(bitmap, 0)
        
        val textBlocks = suspendCancellableCoroutine<List<com.google.mlkit.vision.text.Text.TextBlock>> { continuation ->
            textRecognizer.process(image)
                .addOnSuccessListener { visionText ->
                    continuation.resume(visionText.textBlocks)
                }
                .addOnFailureListener { exception ->
                    Log.e(TAG, "Text recognition failed", exception)
                    continuation.resume(emptyList())
                }
        }
        
        for (block in textBlocks) {
            // 创建文字元素
            val elementRect = Rect(
                block.boundingBox?.left ?: 0,
                block.boundingBox?.top ?: 0,
                block.boundingBox?.right ?: 0,
                block.boundingBox?.bottom ?: 0
            )
            
            // 只添加有实际文本的块
            if (block.text.isNotBlank() && !elementRect.isEmpty) {
                elements.add(
                    VisualElement(
                        id = "text_${elements.size}",
                        text = block.text,
                        bounds = com.autoglm.android.core.models.Rect(
                            left = elementRect.left,
                            top = elementRect.top,
                            right = elementRect.right,
                            bottom = elementRect.bottom
                        ),
                        confidence = 0.95f,
                        type = VisualElementType.TEXT,
                        clickable = guessIfClickable(block.text)
                    )
                )
            }
        }
        
        return@withContext elements
    }
    
    /**
     * 检测UI元素
     */
    private suspend fun detectUIElements(bitmap: Bitmap): List<VisualElement> = withContext(Dispatchers.Default) {
        val elements = mutableListOf<VisualElement>()
        val image = InputImage.fromBitmap(bitmap, 0)
        
        val detectedObjects = suspendCancellableCoroutine<List<com.google.mlkit.vision.objects.DetectedObject>> { continuation ->
            objectDetector.process(image)
                .addOnSuccessListener { results ->
                    continuation.resume(results)
                }
                .addOnFailureListener { exception ->
                    Log.e(TAG, "Object detection failed", exception)
                    continuation.resume(emptyList())
                }
        }
        
        for (detectedObject in detectedObjects) {
            // 根据类别和分数判断是否为UI元素
            val elementType = classifyUIElement(detectedObject)
            
            val elementRect = detectedObject.boundingBox
            
            elements.add(
                VisualElement(
                    id = "ui_${elements.size}",
                    text = "",
                    bounds = com.autoglm.android.core.models.Rect(
                        left = elementRect.left,
                        top = elementRect.top,
                        right = elementRect.right,
                        bottom = elementRect.bottom
                    ),
                    confidence = detectedObject.labels.firstOrNull()?.confidence ?: 0.6f,
                    type = elementType,
                    clickable = elementType != VisualElementType.UNKNOWN
                )
            )
        }
        
        return@withContext elements
    }
    
    /**
     * 使用大模型分析更复杂的界面元素
     */
    private suspend fun analyzeComplexElements(bitmap: Bitmap, detectedElements: List<VisualElement>): List<VisualElement> = withContext(Dispatchers.Default) {
        if (bitmap.width * bitmap.height > 1000000) {
            // 图像太大，需要缩小后再分析
            val scaleFactor = 0.5f
            val resized = Bitmap.createScaledBitmap(
                bitmap,
                (bitmap.width * scaleFactor).toInt(),
                (bitmap.height * scaleFactor).toInt(),
                true
            )
            return@withContext analyzeComplexElements(resized, detectedElements)
        }
        
        try {
            // 将图像转为base64以便发送给LLM
            val base64Image = bitmapToBase64(bitmap)
            
            // 发送给LLM进行分析
            val elementsJson = detectedElements.map { it.toJson() }
            
            val prompt = """
                我是Android应用自动化助手。分析以下屏幕截图，帮我识别更多界面元素，尤其是：
                
                1. 导航栏和标签栏
                2. 按钮和可点击区域
                3. 文本输入框
                4. 开关和选择器
                5. 菜单和弹出框
                
                我已经检测到这些元素：$elementsJson
                
                请分析图像并以JSON格式返回其他你识别到的UI元素。每个元素包含以下字段：
                - id: 元素唯一标识
                - type: 元素类型（BUTTON, INPUT, SWITCH, MENU, TAB, CARD, BANNER, UNKNOWN）
                - bounds: 元素坐标{left, top, right, bottom}
                - text: 元素上的文字（如果有）
                - clickable: 是否可点击
                - confidence: 置信度(0-1)
                
                只返回JSON数组，不要有其他文字。
            """.trimIndent()
            
            val response = llmService.generateWithImage(prompt, base64Image, 0.3)
            
            return@withContext parseElementsFromLLMResponse(response)
        } catch (e: Exception) {
            Log.e(TAG, "Error analyzing complex elements", e)
            return@withContext emptyList()
        }
    }
    
    /**
     * 猜测文本是否可点击
     */
    private fun guessIfClickable(text: String): Boolean {
        // 常见的可点击文本模式
        val clickablePatterns = listOf(
            "确定", "取消", "提交", "保存", "编辑", "删除", "返回", "下一步", "完成",
            "登录", "注册", "搜索", "查询", "设置", "更多", "查看", "关闭", "打开",
            "添加", "新建", "创建", "点击", "立即", "同意", "确认", "Ok", "OK",
            "Yes", "No", "Cancel", "Next", "Back", "Done", "Save", "Submit"
        )
        
        return clickablePatterns.any { pattern ->
            text.contains(pattern, ignoreCase = true)
        }
    }
    
    /**
     * 分类UI元素类型
     */
    private fun classifyUIElement(detectedObject: com.google.mlkit.vision.objects.DetectedObject): VisualElementType {
        val labels = detectedObject.labels
        
        // 检查是否有带有高置信度的分类标签
        for (label in labels) {
            val text = label.text.lowercase()
            val confidence = label.confidence
            
            if (confidence < 0.5f) continue
            
            when {
                text.contains("button") -> return VisualElementType.BUTTON
                text.contains("input") || text.contains("textfield") -> return VisualElementType.INPUT
                text.contains("toggle") || text.contains("switch") -> return VisualElementType.SWITCH
                text.contains("menu") -> return VisualElementType.MENU
                text.contains("tab") -> return VisualElementType.TAB
                text.contains("card") -> return VisualElementType.CARD
                text.contains("banner") -> return VisualElementType.BANNER
            }
        }
        
        // 没有明确分类则返回未知类型
        return VisualElementType.UNKNOWN
    }
    
    /**
     * 将bitmap转为base64字符串
     */
    private fun bitmapToBase64(bitmap: Bitmap): String {
        val outputStream = java.io.ByteArrayOutputStream()
        bitmap.compress(Bitmap.CompressFormat.JPEG, 85, outputStream)
        val byteArray = outputStream.toByteArray()
        return android.util.Base64.encodeToString(byteArray, android.util.Base64.DEFAULT)
    }
    
    /**
     * 从LLM响应中解析元素
     */
    private fun parseElementsFromLLMResponse(response: String): List<VisualElement> {
        val elements = mutableListOf<VisualElement>()
        
        try {
            // 尝试提取JSON数组
            val jsonPattern = "\\[\\s*\\{.*\\}\\s*\\]".toRegex(RegexOption.DOT_MATCHES_ALL)
            val matchResult = jsonPattern.find(response)
            
            val jsonString = matchResult?.value ?: response
            
            val jsonArray = JSONArray(jsonString)
            
            for (i in 0 until jsonArray.length()) {
                val elementJson = jsonArray.getJSONObject(i)
                
                val id = elementJson.optString("id", "llm_element_$i")
                val type = parseElementType(elementJson.optString("type", "UNKNOWN"))
                val text = elementJson.optString("text", "")
                val clickable = elementJson.optBoolean("clickable", false)
                val confidence = elementJson.optDouble("confidence", 0.7).toFloat()
                
                val boundsJson = elementJson.optJSONObject("bounds")
                val bounds = if (boundsJson != null) {
                    com.autoglm.android.core.models.Rect(
                        left = boundsJson.optInt("left", 0),
                        top = boundsJson.optInt("top", 0),
                        right = boundsJson.optInt("right", 0),
                        bottom = boundsJson.optInt("bottom", 0)
                    )
                } else {
                    com.autoglm.android.core.models.Rect(0, 0, 0, 0)
                }
                
                // 只添加坐标有效的元素
                if (!bounds.isEmpty()) {
                    elements.add(
                        VisualElement(
                            id = id,
                            text = text,
                            bounds = bounds,
                            confidence = confidence,
                            type = type,
                            clickable = clickable
                        )
                    )
                }
            }
        } catch (e: Exception) {
            Log.e(TAG, "Error parsing LLM response", e)
        }
        
        return elements
    }
    
    /**
     * 解析元素类型
     */
    private fun parseElementType(typeString: String): VisualElementType {
        return when (typeString.uppercase()) {
            "BUTTON" -> VisualElementType.BUTTON
            "INPUT" -> VisualElementType.INPUT
            "SWITCH" -> VisualElementType.SWITCH
            "MENU" -> VisualElementType.MENU
            "TAB" -> VisualElementType.TAB
            "CARD" -> VisualElementType.CARD
            "BANNER" -> VisualElementType.BANNER
            "TEXT" -> VisualElementType.TEXT
            else -> VisualElementType.UNKNOWN
        }
    }
    
    /**
     * 将VisualElement转换为UIElement
     */
    fun convertToUIElement(visualElement: VisualElement): UIElement {
        return UIElement(
            id = visualElement.id,
            text = visualElement.text,
            contentDescription = visualElement.text,
            resourceId = "",
            className = visualElement.type.name,
            packageName = "",
            bounds = visualElement.bounds,
            isClickable = visualElement.clickable,
            isScrollable = false,
            isLongClickable = visualElement.clickable,
            isVisibleToUser = true,
            isEnabled = true,
            isChecked = null,
            isFocused = null,
            children = emptyList(),
            parent = null
        )
    }
}

/**
 * 视觉元素类型
 */
enum class VisualElementType {
    BUTTON,     // 按钮
    INPUT,      // 输入框
    SWITCH,     // 开关
    MENU,       // 菜单
    TAB,        // 标签
    CARD,       // 卡片
    BANNER,     // 横幅
    TEXT,       // 文本
    UNKNOWN     // 未知类型
} 